[mlir][Vector] Add `vector.shuffle` fold for poison inputs #125608

dcaballe · 2025-02-04T00:36:58Z

#124863 added folding support for poison indices to vector.shuffle. This PR adds support for folding vector.shuffle ops with one or two poison input vectors.

llvmbot · 2025-02-04T00:38:11Z

@llvm/pr-subscribers-mlir-vector

@llvm/pr-subscribers-mlir

Author: Diego Caballero (dcaballe)

Changes

#124863 added folding support for poison indices to vector.shuffle. This PR adds support for folding vector.shuffle ops with one or two poison input vectors.

Full diff: https://github.com/llvm/llvm-project/pull/125608.diff

2 Files Affected:

(modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+32-11)
(modified) mlir/test/Dialect/Vector/canonicalize.mlir (+39)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 93f89eda2da5a6..8d5691f38f273c 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -26,7 +26,6 @@
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/IRMapping.h"
@@ -42,7 +41,6 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
-#include "llvm/ADT/bit.h"
 
 #include <cassert>
 #include <cstdint>
@@ -2696,25 +2694,45 @@ OpFoldResult vector::ShuffleOp::fold(FoldAdaptor adaptor) {
   if (!v1Attr || !v2Attr)
     return {};
 
+  // Fold shuffle poison, poison -> poison.
+  bool isV1Poison = isa<ub::PoisonAttr>(v1Attr);
+  bool isV2Poison = isa<ub::PoisonAttr>(v2Attr);
+  if (isV1Poison && isV2Poison)
+    return ub::PoisonAttr::get(getContext());
+
   // Only support 1-D for now to avoid complicated n-D DenseElementsAttr
   // manipulation.
   if (v1Type.getRank() != 1)
     return {};
 
-  int64_t v1Size = v1Type.getDimSize(0);
+  // Poison input attributes need special handling as they are not
+  // DenseElementsAttr. If an index is poison, we select the first element of
+  // the first non-poison input.
+  SmallVector<Attribute> v1Elements, v2Elements;
+  Attribute poisonElement;
+  if (!isV2Poison) {
+    v2Elements =
+        to_vector(cast<DenseElementsAttr>(v2Attr).getValues<Attribute>());
+    poisonElement = v2Elements[0];
+  }
+  if (!isV1Poison) {
+    v1Elements =
+        to_vector(cast<DenseElementsAttr>(v1Attr).getValues<Attribute>());
+    poisonElement = v1Elements[0];
+  }
 
   SmallVector<Attribute> results;
-  auto v1Elements = cast<DenseElementsAttr>(v1Attr).getValues<Attribute>();
-  auto v2Elements = cast<DenseElementsAttr>(v2Attr).getValues<Attribute>();
+  int64_t v1Size = v1Type.getDimSize(0);
   for (int64_t maskIdx : mask) {
     Attribute indexedElm;
-    // Select v1[0] for poison indices.
     // TODO: Return a partial poison vector when supported by the UB dialect.
     if (maskIdx == ShuffleOp::kPoisonIndex) {
-      indexedElm = v1Elements[0];
+      indexedElm = poisonElement;
     } else {
-      indexedElm =
-          maskIdx < v1Size ? v1Elements[maskIdx] : v2Elements[maskIdx - v1Size];
+      if (maskIdx < v1Size)
+        indexedElm = isV1Poison ? poisonElement : v1Elements[maskIdx];
+      else
+        indexedElm = isV2Poison ? poisonElement : v2Elements[maskIdx - v1Size];
     }
 
     results.push_back(indexedElm);
@@ -3332,13 +3350,15 @@ class InsertStridedSliceConstantFolder final
         !destVector.hasOneUse())
       return failure();
 
-    auto denseDest = llvm::cast<DenseElementsAttr>(vectorDestCst);
-
     TypedValue<VectorType> sourceValue = op.getSource();
     Attribute sourceCst;
     if (!matchPattern(sourceValue, m_Constant(&sourceCst)))
       return failure();
 
+    // TODO: Support poison.
+    if (isa<ub::PoisonAttr>(vectorDestCst) || isa<ub::PoisonAttr>(sourceCst))
+      return failure();
+
     // TODO: Handle non-unit strides when they become available.
     if (op.hasNonUnitStrides())
       return failure();
@@ -3355,6 +3375,7 @@ class InsertStridedSliceConstantFolder final
     // increasing linearized position indices.
     // Because the destination may have higher dimensionality then the slice,
     // we keep track of two overlapping sets of positions and offsets.
+    auto denseDest = llvm::cast<DenseElementsAttr>(vectorDestCst);
     auto denseSlice = llvm::cast<DenseElementsAttr>(sourceCst);
     auto sliceValuesIt = denseSlice.value_begin<Attribute>();
     auto newValues = llvm::to_vector(denseDest.getValues<Attribute>());
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6858f0d56e6412..65c3ab264283d2 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2023,6 +2023,45 @@ func.func @shuffle_1d_poison_idx() -> vector<4xi32> {
 
 // -----
 
+// CHECK-LABEL: func @shuffle_1d_rhs_lhs_poison
+//   CHECK-NOT:   vector.shuffle
+//       CHECK:   %[[V:.+]] = ub.poison : vector<4xi32>
+//       CHECK:   return %[[V]]
+func.func @shuffle_1d_rhs_lhs_poison() -> vector<4xi32> {
+  %v0 = ub.poison : vector<3xi32>
+  %v1 = ub.poison : vector<3xi32>
+  %shuffle = vector.shuffle %v0, %v1 [3, 1, 5, 4] : vector<3xi32>, vector<3xi32>
+  return %shuffle : vector<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @shuffle_1d_lhs_poison
+//   CHECK-NOT:   vector.shuffle
+//       CHECK:   %[[V:.+]] = arith.constant dense<[5, 4, 5, 5]> : vector<4xi32>
+//       CHECK:   return %[[V]]
+func.func @shuffle_1d_lhs_poison() -> vector<4xi32> {
+  %v0 = arith.constant dense<[5, 4, 3]> : vector<3xi32>
+  %v1 = ub.poison : vector<3xi32>
+  %shuffle = vector.shuffle %v0, %v1 [3, 1, 5, 4] : vector<3xi32>, vector<3xi32>
+  return %shuffle : vector<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @shuffle_1d_rhs_poison
+//   CHECK-NOT:   vector.shuffle
+//       CHECK:   %[[V:.+]] = arith.constant dense<[2, 2, 0, 1]> : vector<4xi32>
+//       CHECK:   return %[[V]]
+func.func @shuffle_1d_rhs_poison() -> vector<4xi32> {
+  %v0 = ub.poison : vector<3xi32>
+  %v1 = arith.constant dense<[2, 1, 0]> : vector<3xi32>
+  %shuffle = vector.shuffle %v0, %v1 [3, 1, 5, 4] : vector<3xi32>, vector<3xi32>
+  return %shuffle : vector<4xi32>
+}
+
+// -----
+
 // CHECK-LABEL: func @shuffle_canonicalize_0d
 func.func @shuffle_canonicalize_0d(%v0 : vector<i32>, %v1 : vector<i32>) -> vector<1xi32> {
   // CHECK: vector.broadcast %{{.*}} : vector<i32> to vector<1xi32>

We recently added folding support for poison indices to `vector.shuffle`. This PR adds support for folding poison inputs.

banach-space

LGTM, thanks!

Left a couple of optional nits.

banach-space · 2025-02-04T19:47:33Z

mlir/test/Dialect/Vector/canonicalize.mlir

+  %v0 = arith.constant dense<[5, 4, 3]> : vector<3xi32>
+  %v1 = ub.poison : vector<3xi32>
+  %shuffle = vector.shuffle %v0, %v1 [3, 1, 5, 4] : vector<3xi32>, vector<3xi32>


[nit] There's a value and index == 5, so it's not obvious that the first element of %v0 is in any way significant. Perhaps use some more distinct number? (e.g. 123).

Suggested change

%v0 = arith.constant dense<[5, 4, 3]> : vector<3xi32>

%v1 = ub.poison : vector<3xi32>

%shuffle = vector.shuffle %v0, %v1 [3, 1, 5, 4] : vector<3xi32>, vector<3xi32>

%v0 = arith.constant dense<[123, 4, 3]> : vector<3xi32>

%v1 = ub.poison : vector<3xi32>

%shuffle = vector.shuffle %v0, %v1 [3, 1, 123, 4] : vector<3xi32>, vector<3xi32>

I appreciate that this is obvious right now, but lets also cater for our future selves :)

banach-space · 2025-02-04T19:49:54Z

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

+  // Poison input attributes need special handling as they are not
+  // DenseElementsAttr. If an index is poison, we select the first element of
+  // the first non-poison input.


[nit] To me this is a fairly significant (and not immediately intuitive) part of the design. Perhaps move above the signature?

Also, is this based on some prior-art? Just curious, this does make sense.

Not sure I follow the prior-art part. Do you mean why we pick the first element of the first non-poison input? Poison is basically UB so given that we can't represent a partially poison vector we just make a random decision, which is ok as part of the UB behavior.

It's valid to substitute poison with an arbitrary value

It's valid to substitute poison with an arbitrary value

Sure, but we are selecting a specific "arbitrary value" :)

Not sure I follow the prior-art part.

I was just curious whether there's any rationale behind this specific option. For example, something else in LLVM or MLIR makes similar choice?

Basically, what I'm missing is "why would we select the first element"? Something along the lines would be helpful:

I doesn't matter what we select, but we need to make a choice. We choose the first element.

Sure, but we are selecting a specific "arbitrary value" :)

??? In this context, arbitrary is synonymous to non-deterministics, as in: absolutely any value will do and the choice doesn't have to be fair by any definition of fair.

dcaballe · 2025-02-04T23:26:52Z

It looks like Github has been in "Processing updates" stage (see top of the page) for almost an hour... Weird...

llvm#124863 added folding support for poison indices to `vector.shuffle`. This PR adds support for folding `vector.shuffle` ops with one or two poison input vectors.

dcaballe requested review from Groverkss, banach-space and nicolasvasilache as code owners February 4, 2025 00:36

dcaballe requested a review from kuhar February 4, 2025 00:37

llvmbot added mlir:vectorops mlir mlir:vector labels Feb 4, 2025

kuhar approved these changes Feb 4, 2025

View reviewed changes

[mlir][Vector] Add vector.shuffle fold for poison inputs

46a4887

We recently added folding support for poison indices to `vector.shuffle`. This PR adds support for folding poison inputs.

dcaballe force-pushed the input-poison-shuffle-canon branch from ece1c6d to 46a4887 Compare February 4, 2025 19:34

banach-space approved these changes Feb 4, 2025

View reviewed changes

Feedback

31d1bd3

Empty

6bc3b30

dcaballe merged commit c6eef00 into llvm:main Feb 5, 2025
6 of 7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][Vector] Add `vector.shuffle` fold for poison inputs #125608

[mlir][Vector] Add `vector.shuffle` fold for poison inputs #125608

Uh oh!

dcaballe commented Feb 4, 2025

Uh oh!

llvmbot commented Feb 4, 2025 •

edited

Loading

Uh oh!

banach-space left a comment

Uh oh!

banach-space Feb 4, 2025

Uh oh!

banach-space Feb 4, 2025

Uh oh!

dcaballe Feb 4, 2025

Uh oh!

kuhar Feb 4, 2025

Uh oh!

banach-space Feb 5, 2025

Uh oh!

kuhar Feb 5, 2025 •

edited

Loading

Uh oh!

dcaballe commented Feb 4, 2025 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[mlir][Vector] Add vector.shuffle fold for poison inputs #125608

[mlir][Vector] Add vector.shuffle fold for poison inputs #125608

Uh oh!

Conversation

dcaballe commented Feb 4, 2025

Uh oh!

llvmbot commented Feb 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

banach-space left a comment

Choose a reason for hiding this comment

Uh oh!

banach-space Feb 4, 2025

Choose a reason for hiding this comment

Uh oh!

banach-space Feb 4, 2025

Choose a reason for hiding this comment

Uh oh!

dcaballe Feb 4, 2025

Choose a reason for hiding this comment

Uh oh!

kuhar Feb 4, 2025

Choose a reason for hiding this comment

Uh oh!

banach-space Feb 5, 2025

Choose a reason for hiding this comment

Uh oh!

kuhar Feb 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dcaballe commented Feb 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

[mlir][Vector] Add `vector.shuffle` fold for poison inputs #125608

[mlir][Vector] Add `vector.shuffle` fold for poison inputs #125608

llvmbot commented Feb 4, 2025 •

edited

Loading

kuhar Feb 5, 2025 •

edited

Loading

dcaballe commented Feb 4, 2025 •

edited

Loading